# encoding: utf-8

import json
import pickle
import numpy as np
from sklearn.ensemble import RandomForestRegressor

X, y = [], []

print("load data")
with open("./train.txt", "r", encoding="utf-8") as f:
    lines = f.readlines()

lines = [json.loads(x) for x in lines if x.strip()]
for line in lines:
    X.append(np.array(line[:10]))
    y.append(line[-1])

print("start train")
model = RandomForestRegressor().fit(X, y)

print("train done")

with open("./rfr.pkl", "wb") as f:
    pickle.dump(model, f)
print("save done")

# with open("./rfr.pkl", "rb") as f:
#     model = pickle.load(f)
new_x = np.array([[11.02, 10.06, 9.72, 10.51, 9.99, 9.95, 10.09, 10.33, 9.82, 9.94]])
new_y = 9.53


result = model.predict(new_x)
print(result)  # [9.7212]
